// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <string>

#include <fmt/format.h>
#include <melon/benchmark.h>
#include <melon/conv.h>
#include <melon/random.h>
#include <melon/init/init.h>

#include <pollux/functions/lib/re2_functions.h>
#include <pollux/functions/lib/benchmarks/FunctionBenchmarkBase.h>
#include <pollux/functions/sparksql/in.h>
#include <pollux/vector/vector_fuzzer.h>

namespace kumo::pollux::functions {
void registerPrestoIn() {
  POLLUX_REGISTER_VECTOR_FUNCTION(udf_in, "presto");
}
void registerArrayConstructor() {
  POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_constructor, "array_constructor");
}
} // namespace kumo::pollux::functions

namespace kumo::pollux::functions::sparksql {
namespace {

int in_int(int iters, int inListSize, const std::string& functionName) {
  melon::BenchmarkSuspender kSuspender;
  test::FunctionBenchmarkBase benchmarkBase;

  VectorFuzzer::Options opts;
  opts.vectorSize = 1024;
  auto vector = VectorFuzzer(opts, benchmarkBase.pool()).fuzzFlat(BIGINT());
  auto simpleVector = vector->as<SimpleVector<int64_t>>();
  const auto data = benchmarkBase.maker().row_vector({vector});

  std::string exprStr = functionName + "(c0, array_constructor(";
  for (int i = 0; i < inListSize; i++) {
    if (i > 0) {
      exprStr += ", ";
    }
    exprStr += fmt::format(
        "{}", simpleVector->value_at(melon::Random::rand32() % opts.vectorSize));
  }
  exprStr += "))";
  exec::ExprSet expr = benchmarkBase.compileExpression(exprStr, data->type());
  kSuspender.dismiss();
  for (int i = 0; i != iters; ++i) {
    benchmarkBase.evaluate(expr, data);
  }
  return iters * opts.vectorSize;
}

BENCHMARK_NAMED_PARAM_MULTI(in_int, presto_rhs1, 1, "presto");
BENCHMARK_RELATIVE_NAMED_PARAM_MULTI(in_int, spark_rhs1, 1, "sparkin");
BENCHMARK_NAMED_PARAM_MULTI(in_int, presto_rhs3, 3, "presto");
BENCHMARK_RELATIVE_NAMED_PARAM_MULTI(in_int, spark_rhs3, 3, "sparkin");
BENCHMARK_NAMED_PARAM_MULTI(in_int, presto_rhs10, 10, "presto");
BENCHMARK_RELATIVE_NAMED_PARAM_MULTI(in_int, spark_rhs10, 10, "sparkin");
BENCHMARK_NAMED_PARAM_MULTI(in_int, presto_rhs100, 100, "presto");
BENCHMARK_RELATIVE_NAMED_PARAM_MULTI(in_int, spark_rhs100, 100, "sparkin");
BENCHMARK_NAMED_PARAM_MULTI(in_int, presto_rhs1000, 1000, "presto");
BENCHMARK_RELATIVE_NAMED_PARAM_MULTI(in_int, spark_rhs1000, 1000, "sparkin");

int in_str(int iters, int inListSize, const std::string& functionName) {
  melon::BenchmarkSuspender kSuspender;
  test::FunctionBenchmarkBase benchmarkBase;

  VectorFuzzer::Options opts;
  opts.vectorSize = 1024;
  auto vector = VectorFuzzer(opts, benchmarkBase.pool()).fuzzFlat(VARCHAR());
  auto simpleVector = vector->as<SimpleVector<StringView>>();
  const auto data = benchmarkBase.maker().row_vector({vector});

  std::string exprStr = functionName + "(c0, array_constructor(";
  for (int i = 0; i < inListSize; i++) {
    if (i > 0) {
      exprStr += ", ";
    }
    exprStr += fmt::format(
        "'{}'",
        simpleVector->value_at(melon::Random::rand32() % opts.vectorSize));
  }
  exprStr += "))";
  exec::ExprSet expr = benchmarkBase.compileExpression(exprStr, data->type());
  kSuspender.dismiss();
  for (int i = 0; i != iters; ++i) {
    benchmarkBase.evaluate(expr, data);
  }
  return iters * opts.vectorSize;
}

BENCHMARK_NAMED_PARAM_MULTI(in_str, presto_rhs1, 1, "presto");
BENCHMARK_RELATIVE_NAMED_PARAM_MULTI(in_str, spark_rhs1, 1, "sparkin");
BENCHMARK_NAMED_PARAM_MULTI(in_str, presto_rhs3, 3, "presto");
BENCHMARK_RELATIVE_NAMED_PARAM_MULTI(in_str, spark_rhs3, 3, "sparkin");
BENCHMARK_NAMED_PARAM_MULTI(in_str, presto_rhs10, 10, "presto");
BENCHMARK_RELATIVE_NAMED_PARAM_MULTI(in_str, spark_rhs10, 10, "sparkin");
BENCHMARK_NAMED_PARAM_MULTI(in_str, presto_rhs100, 100, "presto");
BENCHMARK_RELATIVE_NAMED_PARAM_MULTI(in_str, spark_rhs100, 100, "sparkin");
BENCHMARK_NAMED_PARAM_MULTI(in_str, presto_rhs1000, 1000, "presto");
BENCHMARK_RELATIVE_NAMED_PARAM_MULTI(in_str, spark_rhs1000, 1000, "sparkin");

} // namespace

void registerInFunctions() {
  registerIn("spark");
}

} // namespace kumo::pollux::functions::sparksql

int main(int argc, char** argv) {
  melon::Init init{&argc, &argv};
  kumo::pollux::functions::registerPrestoIn();
  kumo::pollux::functions::sparksql::registerInFunctions();
  kumo::pollux::functions::registerArrayConstructor();
  melon::runBenchmarks();
  return 0;
}
